#include"StringPool.h"
#include"variable.h"
#include<vector>
#include<ctime>
#include<iostream>
#include<map>
#include<iomanip>
#include"DependencyInstance.h"
#include"Util.h"
#include"KBestParseForest3OTSibling.h"
using namespace std;

KBestParseForest3OTSibling::KBestParseForest3OTSibling(ParseForestItem ****d_chart, int start, int end, DependencyInstance &inst, int K) : KBestParseForest2OSibling(start, end, inst, K){
	this->d_chart = d_chart;
	for(int i = 0; i < end + 1; i++){
		for(int j = 0; j< end + 1; j++){
			int nn;
			if(i == j){
				nn = 0;
			}
			else if(i > j){
				nn = j;
			}
			else{
				nn = end + 1;
			}
			for(int k = 0; k < nn; k++){
				if(d_chart[i][j][k] != NULL){
					for(int dir = 0; dir < 2; dir++){
						for(int kk = 0; kk < K; kk++){
							int key = getKey(dir, kk);
							d_chart[i][j][k][key] = ParseForestItem(i, j, k, -1, dir, 3, Negative_Infinity, NULL, NULL);
						}
					}
				}
			}
		}
	}
}

KBestParseForest3OTSibling::~KBestParseForest3OTSibling(){}

bool KBestParseForest3OTSibling::add_d(int s, int type, int dir, double score){
	return false;
}

bool KBestParseForest3OTSibling::add_d(int s, int m, int t, int type, int dir, double score, ParseForestItem *p1, ParseForestItem *p2){
	bool added = false;
	int key = getKey(dir, K - 1);
	if(d_chart[s][m][t][key].prob > score){
		return false;
	}
	for(int i = 0;i< K; i++){
		key = getKey(dir, i);
		if(d_chart[s][m][t][key].prob < score){
			ParseForestItem tmp = d_chart[s][m][t][key];
			d_chart[s][m][t][key] = ParseForestItem(s, m, t, type, dir, 3, score, p1, p2);
			for(int j = i + 1; j < K && tmp.prob > Negative_Infinity; j++){
				int _key = getKey(dir, j);
				ParseForestItem tmp1 = d_chart[s][m][t][_key];
				d_chart[s][m][t][_key] = tmp;
				tmp = tmp1;
			}
			added = true;
			break;
		}
	}
	return added;
}

ParseForestItem KBestParseForest3OTSibling::getItem_d(int s, int m, int t, int dir){
	return getItem_d(s, m, t, dir, 0);
}

ParseForestItem KBestParseForest3OTSibling::getItem_d(int s, int m, int t, int dir, int k){
	int key = getKey(dir, k);
	return d_chart[s][m][t][key];
}

ParseForestItem* KBestParseForest3OTSibling::getItems_d(int s, int m, int t, int dir){
	int key = getKey(dir, 0);
	return (d_chart[s][m][t] + key);
}

int KBestParseForest3OTSibling::getKey(int dir, int k){
	return dir * K + k;
}

int KBestParseForest3OTSibling::getKey(int s, int t, int dir, int comp, int k){
	int key = s;
	key = key * (end + 1) + t;
	key = key * 2 + dir;
	key = key * 3 + comp;
	key = key * K + k;
	return key;
}

double KBestParseForest3OTSibling::getProb_d(int s, int m, int t, int dir, int i){
	int key = getKey(dir, i);
	return d_chart[s][m][t][key].prob;
}

double KBestParseForest3OTSibling::getProb_d(int s, int m, int t, int dir){
	return getProb_d(s, m, t, dir, 0);
}

vector<pair<FeatureVector*, string*> >* KBestParseForest3OTSibling::getBestParses(DependencyInstance *inst, DependencyPipe *pipe){
	vector<pair<FeatureVector*, string*> >* d = new vector<pair<FeatureVector*, string*> >(K);
	for(int i = 0; i < K; i++){
		int key = getKey(0, end, 0, 1, i);
		if(chart[key].prob > Negative_Infinity){
			(*d)[i].second = getDepString(chart[key]);
			(*d)[i].first = getFeatureVector(inst, pipe, (*d)[i].second);
		}
		else{
			(*d)[i].first = NULL;
			(*d)[i].second = NULL;
		}
	}
	return d;
}

vector<pair<FeatureVector*, string*> >* KBestParseForest3OTSibling::getBestParses(DependencyInstance *inst, DependencyPipe *pipe, std::vector<vector<vector<double> > > &probs, std::vector<vector<vector<double> > > &probs_trips, std::vector<vector<vector<double> > > &probs_sibs, std::vector<vector<vector<vector<double> > > > &probs_ttrips, std::vector<vector<vector<vector<double> > > > &probs_tsibs, std::vector<vector<vector<vector<double> > > > &nt_probs, bool labeled){
	vector<pair<FeatureVector*, string*> >* d = new vector<pair<FeatureVector*, string*> >(K);
	vector<ParseForestItem> result = vector<ParseForestItem>(K);
	for(int i = 0; i < K; i++){
		result[i] = ParseForestItem(0, -1, end, -1, 0, 1, Negative_Infinity, NULL, NULL);
	}
	for(int i = 1; i < end + 1; i++){
		ParseForestItem *left = getItems(1, i, 1, 1);
		ParseForestItem *right = getItems(i, end, 0, 1);
		vector<BestPair>* pairs = getKBestPairs(left, right);
		int size = (int)(pairs->size());
		bool added = true;
		double prob_0i = probs[0][i][0];
		double nt_prob_s_01 = 0.0;
		double nt_prob_t_00 = 0.0;
		if(labeled){
			nt_prob_s_01 = nt_probs[0][rootType][0][1];
			nt_prob_t_00 = nt_probs[i][rootType][0][0];
			prob_0i += nt_prob_s_01 + nt_prob_t_00;
		}

		double prob_000i = probs_ttrips[0][0][0][i] + probs_tsibs[0][0][i][0] + probs_trips[0][0][i] + probs_sibs[0][i][0];

		for(int k = 0; k < size; k++){
			if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
				break;
			}

			int k1 = (*pairs)[k].first;
			int k2 = (*pairs)[k].second;
			double prob = left[k1].prob + right[k2].prob + prob_0i + prob_000i;
			added = false;
			for(int j = 0; j < K; j++){
				if(prob > result[j].prob){
					ParseForestItem tmp = result[j];
					result[j] = ParseForestItem(0, i, end, -1, 0, 1, prob, &(left[k1]), &(right[k2]));
					for(int m = j + 1; m < K; m++){
						ParseForestItem tmp1 = result[m];
						result[m] = tmp;
						tmp = tmp1;
					}
					added = true;
					break;
				}
			}
			if(!added){
				break;
			}
		}
		delete(pairs);
	}
	for(int i = 0; i < K; i++){
		if(result[i].prob > Negative_Infinity){

			string* str1 = getDepString(*(result[i].left));
			(*d)[i].second = Util::trim(str1);
			pool->push_back(str1);
			int ll = (int)((*d)[i].second->length());
			if(ll != 0){
				(*d)[i].second->append(" ");
			}
			(*d)[i].second->append("0|");
			str1 = Util::intToString(result[i].r);
			str1->append(":");
			(*d)[i].second->append(*str1);
			pool->push_back(str1);
			str1 = Util::intToString(KBestParseForest::rootType);
			str1->append(" ");
			(*d)[i].second->append(*str1);
			pool->push_back(str1);
			str1 = getDepString(*(result[i].right));
			string* str2 = Util::trim(str1);
			pool->push_back(str1);
			(*d)[i].second->append(*str2);
			pool->push_back(str2);
			str1 = (*d)[i].second;
			(*d)[i].second = Util::trim(str1);
			pool->push_back(str1);

			(*d)[i].first = getFeatureVector(inst, pipe, (*d)[i].second);
		}
		else{
			(*d)[i].first = NULL;
			(*d)[i].second = NULL;
		}
	}
	return d;
}

string* KBestParseForest3OTSibling::getDepString(ParseForestItem &pfi){
	if(pfi.left == NULL){
		return pool->get();
	}

	if(pfi.dir == 0 && pfi.comp == 0){
		string* str1 = getDepString(*(pfi.left));
		str1->append(" ");
		string* str2 = getDepString(*(pfi.right));
		str1->append(*str2);
		string* result = Util::trim(str1);
		pool->push_back(str1);
		pool->push_back(str2);
		result->append(" ");
		str1 = Util::intToString(pfi.s);
		result->append(*str1 + "|");
		pool->push_back(str1);
		str1 = Util::intToString(pfi.t);
		result->append(*str1 + ":");
		pool->push_back(str1);
		str1 = Util::intToString(pfi.type);
		result->append(*str1);
		pool->push_back(str1);
		str1 = result;
		result = Util::trim(str1);
		pool->push_back(str1);
		return result;
	}
	else if(pfi.dir == 1 && pfi.comp == 0){
		string* str1 = Util::intToString(pfi.t);
		str1->append("|");
		string* str2 = Util::intToString(pfi.s);
		str2->append(":");
		str1->append(*str2);
		pool->push_back(str2);
		str2 = Util::intToString(pfi.type);
		str2->append(" ");
		str1->append(*str2);
		pool->push_back(str2);
		string* result = str1;
		str1 = getDepString(*(pfi.left));
		str1->append(" ");
		str2 = getDepString(*(pfi.right));
		str1->append(*str2);
		pool->push_back(str2);
		str2 = Util::trim(str1);
		pool->push_back(str1);
		result->append(*str2);
		pool->push_back(str2);
		str1 = result;
		result = Util::trim(str1);
		pool->push_back(str1);
		return result;
	}
	else if(pfi.dir == 0 && pfi.comp == 3){
		string* str1 = getDepString(*(pfi.left));
		string *result = Util::trim(str1);
		pool->push_back(str1);
		int ll = (int)(result->length());
		if(ll != 0){
			result->append(" ");
		}
		str1 = Util::intToString(pfi.s);
		result->append(*str1);
		pool->push_back(str1);
		result->append("|");
		str1 = Util::intToString(pfi.r);
		result->append(*str1);
		pool->push_back(str1);
		result->append(":");
		str1 = Util::intToString(pfi.type);
		result->append(*str1);
		pool->push_back(str1);
		result->append(" ");
		str1 = getDepString(*(pfi.right));
		string *str2 = Util::trim(str1);
		pool->push_back(str1);
		str1 = result;
		str1->append(*str2);
		pool->push_back(str2);
		result = Util::trim(str1);
		pool->push_back(str1);
		return result;
	}
	else if(pfi.dir == 1 && pfi.comp == 3){
		string* str1 = getDepString(*(pfi.left));
		string *result = Util::trim(str1);
		pool->push_back(str1);
		int ll = (int)(result->length());
		if(ll != 0){
			result->append(" ");
		}
		str1 = Util::intToString(pfi.t);
		result->append(*str1);
		pool->push_back(str1);
		result->append("|");
		str1 = Util::intToString(pfi.r);
		result->append(*str1);
		pool->push_back(str1);
		result->append(":");
		str1 = Util::intToString(pfi.type);
		result->append(*str1);
		pool->push_back(str1);
		result->append(" ");
		str1 = getDepString(*(pfi.right));
		string *str2 = Util::trim(str1);
		pool->push_back(str1);
		str1 = result;
		str1->append(*str2);
		pool->push_back(str2);
		result = Util::trim(str1);
		pool->push_back(str1);
		return result;
	}
	else{
		string* str1 = getDepString(*(pfi.left));
		str1->append(" ");
		string* str2 = getDepString(*(pfi.right));
		str1->append(*str2);
		string* result = Util::trim(str1);
		pool->push_back(str1);
		pool->push_back(str2);
		return result;
	}
}

